#pragma once
#include "Common.h"
#include <cstring>

// 单层基数树
//32位寻址空间2^32 / 每页空间2^13 = 基数树直接寻址需要的空间2^16
// BITS = 32/64 - page_shift
template <int BITS>
class RadixTree1
{
public:
	RadixTree1()
	{
		size_t size = sizeof(void*) << BITS; 
		size_t alignedSize = Tool::RoundUp(size); 
		_array = (void**)SystemAlloc(alignedSize >> page_shift);  //
		memset(_array, 0, size);
	}
	// 返回PageID对应的span地址
	void* get(PageID id)
	{
		if (id >> BITS > 0) return nullptr; // 大于2^BITS-1
		return _array;
	}
	void set(PageID id, void* span)
	{
		_array[id] = span;
	}

private:
	//const int length = 1 << BITS;
	void** _array;
};


//2层基数树
template <int BITS>
class RadixTree2
{
public:
	RadixTree2()
	{
		memset(_root, 0, sizeof(_root));
	}
	void* get(PageID id)
	{
		assert(!(id >> BITS)); // id 不能大于2^BITS
		PageID rootId = id >> leaf_bits;
		PageID leafId = id & (leaf_length - 1);

		if (_root[rootId]) return _root[rootId]->values[leafId];
		return nullptr;
	}
	void set(PageID id, void* span)
	{
		assert(id < (1 << BITS));
		PageID rootId = id >> leaf_bits;
		PageID leafId = id & (leaf_length - 1);
		if (_root[rootId] == nullptr) 
		{
			_root[rootId] = new Leaf(); // 按需分配叶节点
			memset(_root[rootId]->values, 0, sizeof(_root[rootId]->values)); // 初始化叶节点数组
		}
		_root[rootId]->values[leafId] = span;
	}
private:

	static const size_t root_bits = 5; // 
	static const size_t root_length = 1 << root_bits; // 2^5
	static const size_t leaf_bits = BITS - root_bits; 
	static const size_t leaf_length = 1 << leaf_bits;

	struct Leaf {
		void* values[leaf_length];
	};
	Leaf* _root[root_length]; 
};

// 3层基数树
template<int BITS>
class RadixTree3
{
public:
	RadixTree3()
	{
		_root = NewNode();
	}
	void* get(PageID id)
	{
		PageID rootId = id >> (leaf_bits + interior_bits);
		PageID interId = (id >> leaf_bits) & (interior_length - 1);
		PageID leafId = id & (leaf_length - 1);
		if ((id >> BITS) || !_root->ptrs[rootId] || !_root->ptrs[rootId]->ptrs[interId])
			return nullptr;
		return reinterpret_cast<Leaf*>(_root->ptrs[rootId]->ptrs[interId])->values[leafId];
	}
	void set(PageID id, void* span)
	{
		assert(!(id >> BITS));
		if (!Ensure(id))  assert(false);
		PageID rootId = id >> (leaf_bits + interior_bits);
		PageID interId = (id >> leaf_bits) & (interior_length - 1);
		PageID leafId = id & (leaf_length - 1);
		reinterpret_cast<Leaf*>(_root->ptrs[rootId]->ptrs[interId])->values[leafId] = span;
	}
private:
	bool Ensure(PageID id)
	{
		PageID rootId = id >> (leaf_bits + interior_bits);
		PageID interId = (id >> leaf_bits) & (interior_length - 1);
		
		if ((rootId >> interior_bits) || (interId >> interior_bits)) return false; // 溢出

		// 中间节点为空，创建
		if (!_root->ptrs[rootId]) _root->ptrs[rootId] = NewNode();
		// 叶节点为空，创建
		if (!_root->ptrs[rootId]->ptrs[interId])
		{
			Leaf* leaf = new Leaf;
			memset(leaf, 0, sizeof(*leaf));
			_root->ptrs[rootId]->ptrs[interId] = reinterpret_cast<Node*>(leaf); // 指向叶子节点
		}
		return true;
	}

	static const size_t interior_bits = (BITS + 2) / 3; // 向上取整
	static const size_t interior_length = 1 << interior_bits;
	static const size_t leaf_bits = BITS - 2 * interior_bits;
	static const size_t leaf_length = 1 << leaf_bits;

	struct Node
	{
		Node* ptrs[interior_length]; // 中间节点
	};
	struct Leaf
	{
		void* values[leaf_length]; // 叶节点
	};
	Node* _root; // 根节点\

	Node* NewNode()
	{
		Node* node = new Node;
		memset(node, 0, sizeof(*node));
		return node;
	}
};